# -*- coding: utf-8 -*-
# @Time    : 2025/1/26
# @Author  : Claude AI
# @File    : fast_least_squares_interpolation.py
# @Comment : 高性能最小二乘插值空洞填充算法

import numpy as np
import cv2
from typing import Tuple, Optional
import warnings


class FastLeastSquaresInterpolation:
    """
    高性能最小二乘插值空洞填充算法
    通过批量处理和向量化计算大幅提升速度
    """

    def __init__(self,
                 max_window_size=9,
                 min_window_size=3,
                 min_valid_points=4,
                 fit_order=1,  # 默认使用平面拟合，速度更快
                 use_weighting=False,  # 默认关闭加权
                 max_holes_batch=1000):
        """
        初始化高性能最小二乘插值器

        Args:
            max_window_size: 最大窗口大小（默认较小）
            min_window_size: 最小窗口大小
            min_valid_points: 拟合所需的最少有效点数
            fit_order: 拟合阶数 (1=平面, 2=二次曲面)
            use_weighting: 是否使用加权最小二乘（默认关闭）
            max_holes_batch: 单批处理的最大空洞数量
        """
        self.max_window_size = max_window_size
        self.min_window_size = min_window_size
        self.min_valid_points = min_valid_points
        self.fit_order = fit_order
        self.use_weighting = use_weighting
        self.max_holes_batch = max_holes_batch

    def fill_disparity_holes_fast(self, disparity_map: np.ndarray) -> np.ndarray:
        """
        高性能空洞填充

        Args:
            disparity_map: 输入视差图 (H, W)

        Returns:
            filled_disparity: 填充后的视差图
        """
        if not isinstance(disparity_map, np.ndarray):
            raise ValueError("输入必须是numpy数组")

        if disparity_map.ndim != 2:
            raise ValueError("输入必须是2D视差图")

        filled_disp = disparity_map.copy().astype(np.float64)

        # 识别空洞位置
        hole_mask = (filled_disp <= 0) | np.isnan(filled_disp) | np.isinf(filled_disp)

        if not np.any(hole_mask):
            return filled_disp.astype(np.float32)

        hole_count = np.sum(hole_mask)
        # print(f"发现 {hole_count} 个空洞，开始高性能填充...")

        # 分批处理空洞
        filled_count = 0
        hole_coords = np.argwhere(hole_mask)

        # 如果空洞太多，优先处理边缘空洞
        if hole_count > self.max_holes_batch:
            # 简单的距离计算（避免昂贵的距离变换）
            edge_holes = self._find_edge_holes(filled_disp, hole_mask, self.max_holes_batch)
            selected_coords = edge_holes
        else:
            selected_coords = hole_coords

        # 批量处理
        batch_size = 100
        for i in range(0, len(selected_coords), batch_size):
            batch_coords = selected_coords[i:i + batch_size]
            batch_filled = self._batch_interpolate_points(filled_disp, batch_coords)

            for (y, x), value in zip(batch_coords, batch_filled):
                if value is not None:
                    filled_disp[y, x] = value
                    filled_count += 1

        # 如果还有剩余空洞，使用快速方法填充
        remaining_holes = (filled_disp <= 0) | np.isnan(filled_disp) | np.isinf(filled_disp)
        if np.any(remaining_holes):
            filled_disp = self._fast_fill_remaining(filled_disp, remaining_holes)
            filled_count += np.sum(remaining_holes)

        # print(f"成功填充 {filled_count}/{hole_count} 个空洞")
        return filled_disp.astype(np.float32)

    def _find_edge_holes(self, disp: np.ndarray, hole_mask: np.ndarray, max_count: int):
        """
        快速找到边缘空洞（避免距离变换）
        """
        # 使用卷积找到空洞边缘
        kernel = np.ones((3, 3), dtype=np.uint8)
        hole_uint8 = hole_mask.astype(np.uint8)

        # 膨胀空洞区域
        dilated = cv2.dilate(hole_uint8, kernel, iterations=1)

        # 边缘空洞 = 膨胀后的空洞 - 原始空洞
        edge_mask = (dilated > hole_uint8) & hole_mask
        edge_coords = np.argwhere(edge_mask)

        # 如果边缘空洞不够，添加一些内部空洞
        if len(edge_coords) < max_count:
            all_coords = np.argwhere(hole_mask)
            remaining = max_count - len(edge_coords)

            # 随机选择一些内部空洞
            internal_mask = hole_mask & ~edge_mask
            internal_coords = np.argwhere(internal_mask)

            if len(internal_coords) > 0:
                indices = np.random.choice(len(internal_coords),
                                        min(remaining, len(internal_coords)),
                                        replace=False)
                additional_coords = internal_coords[indices]
                edge_coords = np.vstack([edge_coords, additional_coords])

        return edge_coords[:max_count]

    def _batch_interpolate_points(self, disp: np.ndarray, coords: np.ndarray) -> list:
        """
        批量插值多个点
        """
        results = []

        for y, x in coords:
            # 使用简化的快速插值
            value = self._fast_interpolate_point(disp, x, y)
            results.append(value)

        return results

    def _fast_interpolate_point(self, disp: np.ndarray, x: int, y: int) -> Optional[float]:
        """
        快速单点插值（简化版本）
        """
        # 使用固定小窗口
        window_size = self.min_window_size
        half_window = window_size // 2

        # 提取局部区域
        x_start = max(0, x - half_window)
        x_end = min(disp.shape[1], x + half_window + 1)
        y_start = max(0, y - half_window)
        y_end = min(disp.shape[0], y + half_window + 1)

        local_region = disp[y_start:y_end, x_start:x_end]

        # 找到有效像素
        valid_mask = (local_region > 0) & ~np.isnan(local_region) & ~np.isinf(local_region)

        if np.sum(valid_mask) < self.min_valid_points:
            return None

        # 使用简单的中值插值（速度很快）
        valid_values = local_region[valid_mask]
        return float(np.median(valid_values))

    def _fast_fill_remaining(self, disp: np.ndarray, hole_mask: np.ndarray) -> np.ndarray:
        """
        快速填充剩余空洞（使用简单方法）
        """
        filled = disp.copy()

        # 方法1：最近邻填充（使用OpenCV）
        hole_uint8 = hole_mask.astype(np.uint8)

        # 使用距离变换进行快速最近邻填充
        try:
            distance = cv2.distanceTransform(hole_uint8, cv2.DIST_L2, 5)
            # 简单的最近邻插值
            filled[hole_mask] = np.mean(disp[~hole_mask])  # 使用全局平均值
        except:
            # 回退到更简单的方法
            filled[hole_mask] = np.median(disp[disp > 0]) if np.any(disp > 0) else 0

        return filled

    def fill_disparity_ultra_fast(self, disparity_map: np.ndarray) -> np.ndarray:
        """
        超快速填充（牺牲一些质量换取速度）
        """
        filled_disp = disparity_map.copy().astype(np.float64)
        hole_mask = (filled_disp <= 0) | np.isnan(filled_disp) | np.isinf(filled_disp)

        if not np.any(hole_mask):
            return filled_disp.astype(np.float32)

        hole_count = np.sum(hole_mask)
        # print(f"超快速模式：填充 {hole_count} 个空洞...")

        # 方法1：使用OpenCV的inpaint（非常快）
        try:
            mask_uint8 = hole_mask.astype(np.uint8) * 255
            # 归一化到0-255范围
            disp_norm = cv2.normalize(filled_disp, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

            # 使用TELEA方法（快速且质量好）
            inpainted = cv2.inpaint(disp_norm, mask_uint8, 3, cv2.INPAINT_TELEA)

            # 恢复原始范围
            filled_disp = cv2.normalize(inpainted.astype(np.float32), None,
                                      filled_disp.min(), filled_disp.max(),
                                      cv2.NORM_MINMAX)

        except:
            # 回退到简单填充
            valid_values = filled_disp[~hole_mask]
            if len(valid_values) > 0:
                filled_disp[hole_mask] = np.median(valid_values)

        return filled_disp.astype(np.float32)


def fill_disparity_fast(disparity_map: np.ndarray, mode: str = 'fast') -> np.ndarray:
    """
    便捷函数：高性能填充视差图空洞

    Args:
        disparity_map: 输入视差图
        mode: 模式 ('ultra_fast', 'fast', 'balanced')

    Returns:
        filled_disparity: 填充后的视差图
    """
    if mode == 'ultra_fast':
        # 超快速模式（主要用于实时应用）
        interpolator = FastLeastSquaresInterpolation(
            max_window_size=3,
            min_window_size=3,
            fit_order=1,
            use_weighting=False
        )
        return interpolator.fill_disparity_ultra_fast(disparity_map)

    elif mode == 'fast':
        # 快速模式
        interpolator = FastLeastSquaresInterpolation(
            max_window_size=7,
            min_window_size=3,
            fit_order=1,
            use_weighting=False,
            max_holes_batch=500
        )
        return interpolator.fill_disparity_holes_fast(disparity_map)

    else:  # balanced
        # 平衡模式（质量和速度的折中）
        interpolator = FastLeastSquaresInterpolation(
            max_window_size=9,
            min_window_size=5,
            fit_order=1,
            use_weighting=True,
            max_holes_batch=800
        )
        return interpolator.fill_disparity_holes_fast(disparity_map)


# 兼容性函数
def fill_disparity_least_squares_fast(disparity_map: np.ndarray, mode='fast', **kwargs):
    """
    快速版本的兼容函数
    """
    return fill_disparity_fast(disparity_map, mode)


def fill_disparity_morphology(disp):
    disp_copy = disp.copy()
    invalid = disp <= 0

    # 闭运算填小孔（先膨胀再腐蚀）
    disp_morph = cv2.morphologyEx(disp_copy, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))

    # 替换无效区域
    disp_copy[invalid] = disp_morph[invalid]
    return disp_copy


    

if __name__ == "__main__":
    # 快速测试
    print("高性能最小二乘插值模块测试")

    # 创建测试数据
    test_disp = np.random.rand(100, 100) * 50 + 10
    test_disp[20:30, 40:50] = 0
    test_disp[60:70, 70:80] = 0

    # 测试不同模式
    modes = ['ultra_fast', 'fast', 'balanced']

    for mode in modes:
        print(f"\n测试模式: {mode}")
        import time
        start_time = time.time()
        # 应用填充
        for _i in range(1000):
            filled = fill_disparity_fast(test_disp.copy(), mode=mode)
            # filled = fill_disparity_morphology(test_disp.copy())

        end_time = time.time()
        processing_time = end_time - start_time

        holes = np.sum((filled <= 0) | np.isnan(filled))
        print(f"  处理时间: {processing_time:.3f}s")
        print(f"  剩余空洞: {holes}")

    print("\n测试完成")